#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
心理陪伴智能体 (状态管理依赖MemorySystem)
"""
from datetime import datetime, timedelta
# import datetime
import sys, os
import logging
import json
import uuid
from typing import Dict, List, Any, Optional
from .base_agent import BaseAgent
try:
    from utils.llm_service import LLMService
    from utils.memory_system import MemorySystem
except ImportError as e:
    logging.error(f"导入依赖失败: {e}")
    LLMService = None
    MemorySystem = None

logger = logging.getLogger("AI-MindCare-System-Companion")
# 清除原有可能使用错误编码的handler
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 添加UTF-8编码的handler
handler = logging.StreamHandler(stream=sys.stdout)
handler.setFormatter(
    logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                      datefmt='%Y-%m-%d %H:%M:%S'))
handler.encoding = 'utf-8'
logger.addHandler(handler)


class CompanionAgent(BaseAgent):
    """心理陪伴智能体"""

    def __init__(
            self,
            name: str = "心理陪伴智能体",
            specialty: str = "general",  # 专业
            memory_system: Optional[MemorySystem] = None,
            llm_service: Optional[LLMService] = None):
        super().__init__(name=name,
                         role="companion",
                         memory_system=memory_system)
        self.specialty = specialty
        self.llm_service = llm_service
        if not self.llm_service: logger.error(f"心理陪伴智能体 {name}: LLM 服务未提供。")
        if not self.memory_system:
            logger.error(f"心理陪伴智能体 {name}: MemorySystem 未提供。")
        logger.info(f"心理陪伴智能体 {name} ({specialty}) 初始化完成")

    def _handle_casual_chat_counseling(
            self, context: Dict[str, Any]) -> Dict[str, Any]:
        long_term_memory = context.get("long_term_memory", {})
        # 如果 long_term_memory 是字典，将其格式化为漂亮的字符串
        if isinstance(long_term_memory, dict):
            long_term_memory_str = json.dumps(long_term_memory,
                                              ensure_ascii=False,
                                              indent=2)
        else:
            long_term_memory_str = str(long_term_memory)
        prompt = f"""作为一个AI心理伙伴，请根据以下对话历史、用户的记忆信息和最新消息，生成一个适当的回应。
你既能提供日常友好的聊天，也能在需要时提供专业的心理支持。

结合上下文信息:
{json.dumps(context.get('messages', []), ensure_ascii=False, indent=2)}
判定是否足够充分回答用户的提问，若不能，请参照
用户的长期记忆:
{long_term_memory_str}，结合以下要求进行回答。

请根据用户消息的具体内容和情感需求，自动调整你的回应方式：

一、当用户在进行日常交流时，你会：
1. 保持对话的自然和友好
2. 体现积极倾听和理解
3. 适当引用用户的记忆信息，增强对话的个性化
4. 表达温暖和支持，建立信任关系

二、当用户表达心理困扰或需要专业支持时，你会：
1. 把用户当作具有内在智慧和恢复力的个体对待，尊重用户的直觉、经验与独特性
2. 带着深度理解和共情力，陪伴用户面对困扰
3. 每次生成一个专业、有针对性的问题，进一步了解用户情况以帮助用户理解情绪困扰背后的意义与未满足的需求
4. 基于心理学的研究与实践经验，给出稳定、可靠、逻辑清晰的分析与建议
5. 用系统思维帮助用户整理"碎片化的情绪、反应、经历"，构建清晰的认知结构
6. 结合记忆信息，提供高度个性化的支持，增强用户被理解的感受

三、你在任何情况下都不会：
1.一次性抛出多个问题，造成用户认知负担或对话混乱。
2. 在信息不足时提供草率建议，确保建议基于充分理解。
3. 在无新信息时随意改变立场，导致用户困惑。
4. 用通用化建议掩盖用户独特的心理需求。

四、紧急情况处理指南（高度警惕以下危机信号）：
1. 自伤或自杀的表述或暗示
2. 极度抑郁、绝望或无助的表达
3. 严重的精神症状描述（如幻觉、妄想、极度情绪波动等）
4. 暴力倾向或伤害他人的想法

当察觉到以上任何一种情况，你必须：
1. 保持冷静，表达真诚关切（如"我很担心你现在的状态"）。
2. 明确建议用户立即寻求专业医疗帮助，强调这是必要且关键的步骤。
3. 提供具体求助渠道（如拨打当地心理危机热线）。
4. 清晰说明此类问题超出AI支持范围，需专业医疗介入。

五、心理支持的专业框架
你的支持应逐渐遵循以下结构化步骤：
1. 情绪确认与情境分析：明确用户的情绪（如"愤怒""无助"）及触发事件，确认"发生了什么"和"感受是什么"。
2. 心理机制探索：通过提问引导，深入分析情绪或行为背后的心理模式、认知结构或成长经历（如自动化负面思维、童年影响）。
3. 策略性支持：在获得足够多信息，明确情绪与心理机制后，基于整合模型（认知-情绪-行为-关系-存在）提出具体建议，如认知重构、情绪调节技巧或关系沟通策略。
4. 动态评估与调整：根据用户反馈，持续观察其情绪变化，适时调整支持方式，确保建议有效且贴合需求。
5. 如长期记忆和上下文信息出现矛盾、长期记忆存在矛盾，以上下文信息>最新长期记忆>较旧记忆的顺序进行取舍。
"""
        try:
            response = self.llm_service.generate_response(
                prompt=prompt,
                system_message="你是一个专业的AI心理伙伴，既能提供温暖友好的日常交流，也能提供专业的心理支持和咨询。",
                temperature=0.5,
                max_tokens=500)
            return response
        except ValueError as e:
            logger.error("生成回复时出错: %s", str(e))
            return "您好，非常抱歉，我好像没有理解您的意思呢。为了更好地帮助您，您能否换一种表达方式，或者提供更多的背景信息呢？感谢您的配合！"  # Fallback


#-------------------------------------------------

    def _process_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """处理接收到的消息（通常来自Orchestrator）编排器"""
        content = message.get("content", {})
        sender_id = message.get("sender_id", "unknown")

        if not self.llm_service or not self.memory_system:
            logger.error(f"陪伴助手 {self.name} 无法处理请求，缺少 LLM 或 MemorySystem。")
            return self.send_message(sender_id, {
                "status": "error",
                "message": "内部服务错误，无法处理诊断请求。"
            })

        if "conversation_request" in content:  # 首次回复
            return self._handle_conversation_request(message)
        elif "user_response" in content:  # 非首次回复
            return self._handle_user_response(message)
        else:
            logger.warning(f"{self.name} 收到未知类型的消息: {list(content.keys())}")
            return self.send_message(sender_id, {
                "status": "unhandled",
                "message": f"无法处理此请求类型。"
            })

    # --- 内部处理方法 ---

    def _handle_conversation_request(
            self, message: Dict[str, Any]) -> Dict[str, Any]:
        """处理请求，启动交互式咨询"""

        content = message.get("content", {})

        request = content.get("conversation_request", {})
        user_message = request.get("raw_message", "")
        user_id = request.get("user_id")
        # --- 关键修改：获取 consultation_id ---
        conversation_id = request.get("conversation_id")
        # ------------------------------------
        sender_id = message.get("sender_id", "unknown")

        if not user_id or not conversation_id:  # 确保两个ID都存在
            logger.error("诊断请求缺少 user_id 或 conversation_id")
            return self.send_message(sender_id, {
                "status": "error",
                "message": "诊断请求缺少必要ID"
            })

        logger.info(
            f"{self.name} 开始处理用户 {user_id} 的请求 (Conversation: {conversation_id})。"
        )

        initial_context = {
            "user_id":
            user_id,  # 确保 user_id 正确传递
            "messages": [{
                "role": "user",
                "content": user_message,
                "timestamp": datetime.now().isoformat()
            }],
            "topic":
            "",  # 本次会话的主题
            "create_time":
            datetime.now().isoformat(),  # 创建时间
            "last_updated":
            datetime.now().isoformat(),  # 最近更新时间
        }

        try:
            first_question = self._handle_casual_chat_counseling(  # 使用统一的回复提示词
                message)
            initial_context["messages"].append({
                "role":
                "companion",
                "content":
                first_question,
                "timestamp":
                datetime.now().isoformat()
            })

            # --- 关键修改：使用 conversation_id 作为 key 保存上下文 ---
            self.memory_system.save_conversation_context(
                conversation_id, initial_context)
            # -----------------------------------------------

            response_content = {
                "status": "in_conversation",
                "requires_user_input": True,
                "question": first_question,
                # --- 关键修改：传递正确的 conversation_id 作为 context_id ---
                "context_id": conversation_id,
                # ----------------------------------------------------
                "user_id": user_id
            }
            return self.send_message(sender_id, response_content)

        except Exception as e:
            logger.error(
                f"为用户 {user_id} 生成首个问题时出错 (Consultation: {consultation_id}): {e}",
                exc_info=True)
            return self.send_message(sender_id, {
                "status": "error",
                "message": f"开始服务时出错: {e}"
            })

    def _handle_user_response(self, message: Dict[str, Any]) -> Dict[str, Any]:

        content = message.get("content", {})
        response_data = content.get("user_response", {})
        user_id = message.get("user_id")
        answer = response_data.get("answer", "")

        context_id = message.get("conversation_id")
        logger.info(f"收到来自 user {user_id} 的回答 (Context: {context_id})")
        # ----------------------------------------------------
        sender_id = message.get("sender_id", "unknown")

        if not context_id or not user_id:
            logger.error("用户回复消息缺少 context_id 或 user_id")
            return self.send_message(sender_id, {
                "status": "error",
                "message": "用户回复缺少必要ID"
            })
        if not answer:
            logger.warning(f"收到来自 user {user_id} 的空回复 (Context: {context_id})")
            # 可以选择追问或忽略，这里简单返回错误提示
            return self.send_message(sender_id, {
                "status": "error",
                "message": "收到的回复内容为空，请重新输入。"
            })

        current_context = self.memory_system.get_conversation_context(
            user_id, context_id)

        logger.info(
            f"{self.name} 收到用户 {user_id} 对问题 (Context: {context_id}) 的回答: '{answer[:50]}...'"
        )

        try:
            current_context["messages"].append({
                "role":
                "user",
                "content":
                answer,
                "timestamp":
                datetime.now().isoformat()
            })

            next_question = self._handle_casual_chat_counseling(  # 使用统一的回复提示词
                message)
            current_context["messages"].append({
                "role":
                "companion",
                "content":
                next_question,
                "timestamp":
                datetime.now().isoformat()
            })
            current_context["last_updated"] = datetime.now().isoformat()

            self.memory_system.update_conversation_context(
                user_id, context_id, current_context)  # 更新上下文

            if current_context["topic"] == "" and len(
                    current_context["messages"]) >= 5:
                logger.info(f"为用户 {user_id} 生成主题 (Conversation: {context_id})")
                self.memory_system.generate_topic_for_conversation(
                    user_id, context_id)

            response_content = {
                "status": "in_conversation",
                "requires_user_input": True,
                "question": next_question,
                "context_id": context_id,  # 继续传递
                "user_id": user_id
            }
            return self.send_message(sender_id, response_content)

        except Exception as e:
            logger.error(
                f"为用户 {user_id} 生成下一个问题时出错 (Conversation: {context_id}): {e}",
                exc_info=True)
            return self.send_message(
                sender_id, {
                    "status":
                    "error",
                    "message":
                    f"为用户 {user_id} 生成下一个问题时出错 (Conversation: {context_id}): {e}",
                })
